import tensorflow as tf
import os
import numpy as np
from COMP7015_Mini_Project_2.Model import network
from COMP7015_Mini_Project_2.DataLoader import Cifar10_DataLoader

project_dir = 'drive/My Drive/Colab Notebooks/HKBU_AI_Classs/COMP7015_Mini_Project_2/CKPT'
project_dir_npy = 'drive/My Drive/Colab Notebooks/HKBU_AI_Classs/COMP7015_Mini_Project_2/NPY'
mini_learn_rate = 1e-7
presist_period = 5

def get_config():
    gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.8)
    config = tf.ConfigProto(allow_soft_placement=True, gpu_options=gpu_options)
    config.gpu_options.allow_growth = True
    return config

def move_avg(record, data):
    if len(record)==0:
        record.append(data)
    else:
        value = record[-1]*0.98 + data*0.02
        record.append(value)
    return record[-1]

class Training_Engine:
    def train(self, model, sess, is_restore, epoches, dataLoader, batch_size, is_save, label):

        test_images, test_labels = dataLoader.get_test_batch(2048, label)
        period = 0
        learn_rate = 0.001

        sess.run(tf.global_variables_initializer())

        if is_restore:
            model.saver.restore(sess, save_path=os.path.join(project_dir, '%s.ckpt' % model.scope))

        train_accuracy_record, train_cost_record = [], []
        test_accuracy_record, test_cost_record = [], []

        batch_num_per_epoch = dataLoader.get_train_count()//batch_size
        for i in range(epoches * batch_num_per_epoch):
            images_batch, labels_batch = dataLoader.get_train_batch(batch_size, label)
            _, acc, err = sess.run(fetches=[model.grad_updates, model.accuracy_train, model.logits_loss_train],
                                                  feed_dict={
                                                      model.inputs_train: images_batch,
                                                      model.labels_train: labels_batch,
                                                      model.learn_rate: learn_rate
                                                  })

            last_acc = move_avg(train_accuracy_record, acc)
            last_err = move_avg(train_cost_record, err)
            print('eposch: %d' % i, ' accuracy: ', last_acc, ' loss: ', last_err)

            if last_acc > 0.98:
                print('train acc has completed')
                break

            if (i + 1) % 20 == 0:
                test_acc, test_err = sess.run(fetches=[model.accuracy_test, model.logits_loss_test],
                                              feed_dict={
                                                  model.inputs_test: test_images,
                                                  model.labels_test: test_labels,
                                              })

                last_test_acc = move_avg(test_accuracy_record, test_acc)
                last_test_err = move_avg(test_cost_record, test_err)

                period += 1
                if len(test_cost_record) > 1:
                    if test_cost_record[-2] < test_cost_record[-1] and test_acc>0.5 and period>presist_period:
                        if learn_rate > mini_learn_rate:
                            learn_rate = learn_rate * 0.1
                            period = 0
                        else:
                            print('mini learn rate has complete')
                            break

                    print('test accuracy: ', last_test_acc, 'test cost: ', last_test_err,
                          ' cost decrease:', (test_cost_record[-2] - test_cost_record[-1]),
                          ' learn rate:', learn_rate)

            if (i + 1) % 100 == 0 and is_save:
                model.saver.save(sess, save_path=os.path.join(project_dir, '%s.ckpt' % model.scope),
                           write_state=False, write_meta_graph=False)
        if is_save:
            model.saver.save(sess, save_path=os.path.join(project_dir, '%s.ckpt' % model.scope),
                       write_state=False, write_meta_graph=False)

        save_dict = {
            'train_acc': np.array(train_accuracy_record),
            'train_cost': np.array(train_cost_record),
            'test_acc': np.array(test_accuracy_record),
            'test_cost': np.array(test_cost_record),
        }
        np.save(os.path.join(project_dir_npy, '%s.npy' % model.scope), save_dict)


def run(label, epoches, is_save):
    dataLoader = Cifar10_DataLoader()
    model = network(scope='label%d'%label, num_classes=2)
    engine = Training_Engine()

    with tf.Session(config=get_config()) as sess:
        engine.train(model, sess, is_restore=False, epoches=epoches, dataLoader=dataLoader, batch_size=64, is_save=is_save, label=label)

if __name__ == '__main__':
    run(label=0, epoches=20, is_save=True)
